#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <math.h>
#include <time.h>

// uncomment these lines to compile on a machine //
// that does not support doubles //
// WARNING: single precision will most likely //
// cause serious rounding errors //
//
// #define double float
// #define sqrt sqrtf

// DEFAULT PARAMETERS //

// Effective step size //
const double step_size = 1e-4;
__constant__ double _step_size;

// Unit of length sigma, used to calculate r_c (atomic diameter) //
const double sigma = 1;
__constant__ double _sigma;

// Spring constant for FENE bond //
const double k_bond = 200;
__constant__ double _k_bond;

// Spring constant for harmonic spring (bending) //
// k_bend should equal 10*epsilon/(sigma^2) //
const double k_bend = 20;
__constant__ double _k_bend;



// USER PROMPTED PARAMETERS //

/// @todo make these parameters settable by the user

// Total number of beads //
// Prompted from user //
const int num_beads = 6144;
__constant__ __device__ int _num_beads;

// Number of steps //
// Prompted by user //
const int steps = 1000;
__constant__ int _steps;

// w_c is usually in terms of sigma, so define: w_c = mult * sigma //
// Primary tuning parameter for phase-space //
const double w_c_mult = 1.5;
__constant__ double _w_c_mult;

// Unit of energy epsilon //
// Primary tuning parameter (with w_c_multiplier) for phase space //
/// @todo: change to include temperature as parameter, e.g. eps = k_b * T
const double epsilon = 1.0;
__constant__ double _epsilon;

const double r_c_mod = pow(2, 1.0/6);
__constant__ double _r_c_mod;



// CALCULATED PARAMETERS //

const int num_lipids = num_beads / 6;

// Dimensions of box; calculated to ensure box-spanning bilayer based on # of beads //
//Add 1 to take into account translation inside of walls //
const double boxX = (sqrt(num_lipids)+1) * r_c_mod * 0.95 * sigma;;
__constant__ double _boxX;
const double boxY = boxX;
__constant__ double _boxY;
const double boxZ = 15 * boxX;
__constant__ double _boxZ;

// Range of attractive potential, a primary tuning parameter //
const double w_c = w_c_mult * sigma;
__constant__ double _w_c;




/**
 * @brief Inits the constants to be used on the GPU
 *
 * This function reads constant values from the host
 * and copies them to the constant memory on the GPU
 *
 */
void init_params()
{
	cudaMemcpy(&_boxX, &boxX, sizeof(double), cudaMemcpyHostToDevice);
	cudaMemcpy(&_boxY, &boxY, sizeof(double), cudaMemcpyHostToDevice);
	cudaMemcpy(&_boxZ, &boxZ, sizeof(double), cudaMemcpyHostToDevice);
	cudaMemcpy(&_w_c_mult, &w_c_mult, sizeof(double), cudaMemcpyHostToDevice);
	cudaMemcpy(&_epsilon, &epsilon, sizeof(double), cudaMemcpyHostToDevice);
	cudaMemcpy(&_sigma, &sigma, sizeof(double), cudaMemcpyHostToDevice);
	cudaMemcpy(&_step_size, &step_size, sizeof(double), cudaMemcpyHostToDevice);
}



/**
 * @brief Sets the initial position of the beads
 *
 * Assigns initial positions of beads in a n by n grid.
 * Assumes that (n*8)^2*6 beads are used.
 *
 * @todo This function should be ported to the GPU,
 * 	so that there is some speed up on the initialization
 * 	of huge systems.
 *
 * @todo A switch function should be made available,
 * 	so that the it is possible to load an initial status
 * 	taken from empirical observations.
 */
void init_positions(double* bead_x, double* bead_y, double* bead_z) {
	int flag = 0; //alternates between 0/1 to create lipids facing in opposite directions
	//define spacing parameter for beads = head bead diameter.
	//tail beads will quickly move apart as necessary.
	double r_c = r_c_mod * 0.95 * sigma;
	int xCount = 0; //corresponds to xth most position in grid (if looking at x-z plane)
	int yCount = 0; //yth most position in x-z plane
	//Define dimension of square grid, i.e. number of bead "columns" per side
	int dim = ceil(sqrt(num_beads/6))-1;
	//Translate 1 radius from box walls
	double transX = r_c / 2; //length
	double transY = r_c / 2; //width
	double transZ = boxZ / 2; //depth; roughly center
	//Places 3 beads at once (1 lipid) so make sure numbeads % 3 = 0
	for (int i=0; i < num_beads; i+=3) {
		// Columns should be tangential so that columns span box; hence increment
		// by r_c (e.g. 1 diameter). Note that all x/y coordinates are same for
		// beads in a given column (only z values differ)
		bead_x[i] = transX + xCount * r_c;
		bead_x[i+1] = transX + xCount * r_c;
		bead_x[i+2] = transX + xCount * r_c;
		//y coord; increment after each row of x coords is done
		bead_y[i] = transY + yCount * r_c;
		bead_y[i+1] = transY + yCount * r_c;
		bead_y[i+2] = transY + yCount * r_c;
		// z coord; depends on value of flag (to generate bilayer, alternate flag)
		// creates columns of 6 beads in z direction, with tails facing each other
		if (flag) {
			//Place such that there is 1 r_c space between top/bottom lipid
			bead_z[i] = transZ + 6*r_c;
			bead_z[i+1] = transZ + 5*r_c;
			bead_z[i+2] = transZ + 4*r_c;
			flag = 0;
		} else {
			bead_z[i] = transZ;
			bead_z[i+1] = transZ + r_c;
			bead_z[i+2] = transZ + 2*r_c;
			flag = 1;
		}
		//Want to add such that square grid is formed, manipulate x/z counts accordingly
		if (flag == 0 && xCount < dim) {
			xCount++;
		} else if (flag == 0 && xCount >= dim) {
			xCount = 0;
			yCount++;
		}
	}
}



/**
 * @brief Calculates the distance between two atoms
 *
 * This version corrects the distance
 * by taking into consideration space boundaries
 */
__device__ inline void calc_dist(double* bead_w, int j, int k, double* dist_w)
{
	*dist_w = bead_w[j] - bead_w[k];
	if (abs(*dist_w) > _boxX/2)
		//Correct based on direction (since vector direction matters from above)
		if (*dist_w < 0)
			*dist_w = _boxX + *dist_w;
		else
			*dist_w = *dist_w - _boxX;
}



/**
 * @brief Calculates the distance between two atoms
 *
 * This version does NOT correct the distance
 * by taking into consideration space boundaries
 */
__device__ inline void calc_dist_(double* bead_w, int j, int k, double* dist_w)
{
	*dist_w = bead_w[j] - bead_w[k];
}



/**
 * @brief Calculate WCA force constant based on given atomic distance.
 *
 * Constant given as: Fc = 4*epsilon/r^2 * (12*(b/r)^12 - 6*(b/r)^6)
 * Decompose to each cartesian component by: Fx = Fc*(delta_X) and etc.
 *
 * b is ~ sigma, but < sigma for head beads
 *
 * Note that this force defines r_c as atomic diameter; that is, force acts
 * to repel two beads that have centers within 1 r_c of each other, so r_c is
 * equal to two radii = diameter.
 *
 * @author Reid Van Lehn
 * @warning removed sign, seems stable now. Is this right?
 */
__device__ inline double calcWCA(double b, double dist) {
	double c = b/dist;
	if (pow(dist, 6) > 2*pow(b, 6))
		return 0;
	else
		return 4*_epsilon/(dist*dist)*(12*pow(c, 12) - 6*pow(c, 6));
}



/**
 * @brief Calculate Bending force constant based on given atomic distance.
 *
 * Constant given as: Fc = k_bend * (4*sigma/r - 1)
 *
 * @author Reid Van Lehn
 *
 */
__device__ inline double calcBending(double dist) {
	return k_bend*(4*_sigma/dist - 1);
}



/**
 * @brief Calculate FENE force constant
 *
 * Constant given as:
 * Fc =k_bond * (sigma/r - 1)
 * (< 0 if distance > sigma)
 *
 * @author Reid Van Lehn
 */
__device__ inline double calcFENE(double dist) {
	return k_bond*(_sigma/dist - 1);
}



/**
 * @brief Calculates constant for attractive force
 *
 * Constant given as:
 * 0 for r < rc, > rc + wc
 *
 * @author Reid Van Lehn
 */
__device__ inline double calcAttr(double b, double dist) {
	double r_c = _r_c_mod*b;
	if (dist < r_c || dist > (r_c + w_c))
		return 0;
	else
		return (-M_PI*_epsilon/(dist*w_c))*cos(M_PI*(dist-r_c)/(2*w_c))*sin(M_PI*(dist-r_c)/(2*w_c));
}



/**
 * @brief Computes all the force constants for a pair of beads
 *
 * This function decides which interactions are working
 * between each given couple of beads, and computes the
 * constants for the applicable forces
 *
 * @warning There is no cutoff distance for the interactions
 *	to allow for the consideration of long-range electrostatic
 *	forces
 */
__device__ inline void calc_forces(int i, int j, double* bead_x, double* bead_y, double* bead_z, double* force_x, double* force_y, double* force_z)
{
	double xDist, yDist, zDist, dist; //distances in each component direction + magnitude
	double sumX, sumY, sumZ; //sum of forces acting on current pair
	//Force constants, then decomposed to component directions
	double sumConst, wcaConst, bendConst, feneConst, attrConst;
	sumConst = wcaConst = bendConst = feneConst = attrConst = 0;
	
	// These distances are eventually used to calculate positions when multiplied
	// by force vectors, so sign matters. Take distance vector pointing from 2 to 1
	calc_dist(bead_x, i, j, &xDist);
	calc_dist(bead_y, i, j, &yDist);
	calc_dist_(bead_z, i, j, &zDist);
	
	//Then actually compute the magnitude of distance vector (always positive)
	dist = sqrt(xDist*xDist+yDist*yDist+zDist*zDist);
	
	//Determine WCA force between all i, j beads, broken down into components
	//Note that value of 'b' depends on type of interacting particles
	// All head bead interactions use b = 0.95 sigma
	wcaConst = calcWCA( (!(i % 3) || !(j % 3)) ? (0.95 * _sigma) : (_sigma), dist);
	
	//Every 3rd bead by default is a head bead
	if (!(i % 3) && j == i+2)
		//apply bending force only to 2nd tail bead for this head
		bendConst = calcBending(dist);
	
	if (i % 3 != 2 && j == i+1)
		//apply FENE bonds if i is not the 2nd tail bead
		feneConst = calcFENE(dist);
	
	//No attractive interaction between tail beads in same lipid
	if (i % 3 && j % 3 && j > i+1)
		//apply attractive potential between non-head beads
		attrConst = calcAttr(_sigma, dist);
	
	//Sum constants then decompose to cartesian coordinates
	sumConst = wcaConst + bendConst + feneConst + attrConst;
	//transform to cartesian coords
	sumX = sumConst * xDist;
	sumY = sumConst * yDist;
	sumZ = sumConst * zDist;
	// Put these as equal and opposite forces for particle i and j;
	// then add to current forces. Note: must reset force array in
	// between every step
	__syncthreads();
	force_x[i] += sumX;
	force_x[j] -= sumX;
	force_y[i] += sumY;
	force_y[j] -= sumY;
	force_z[i] += sumZ;
	force_z[j] -= sumZ;
}


/**
 * @brief Takes care of computing the force constants for all beads
 *
 * @note This function only works for the first half of the beads,
 *	as it calculates the interactions between a bead and the n/2 consecutive ones
 */
__global__ void calc_forces_1st(long num_beads, double* bead_x, double* bead_y, double* bead_z, double* force_x, double* force_y, double* force_z) {
	int id = blockDim.x * blockIdx.x + threadIdx.x;
	int its = num_beads >> 1;
	for(int i=1; i<=its; ++i)
	{
		calc_forces(id, id+i, bead_x, bead_y, bead_z, force_x, force_y, force_z);
	}
}



/**
 * @brief Takes care of computing the force constants for all beads
 *
 * @note This function only works for the second half of the beads,
 *	as it calculates the interactions between a bead and the n/2-1
 *	consecutive ones, and corrects the ids so that the first items
 *	in the vector can be accessed when all the items from the second
 *	part are evaluated
 *
 */
__global__ void calc_forces_2nd(long num_beads, double* bead_x, double* bead_y, double* bead_z, double* force_x, double* force_y, double* force_z) {
	int its = num_beads >> 1;
	int id = blockDim.x * blockIdx.x + threadIdx.x + its;
	for(int i=1; i<its; ++i)
	{
		// loop back to the beginning of the list
		int cmp = id+i;
		if(cmp >= num_beads) cmp -= num_beads;
		calc_forces(id, cmp, bead_x, bead_y, bead_z, force_x, force_y, force_z);
	}
}



/**
 * @brief Ensures that an atom is within the given box size
 *
 * This function checks if an atom is within the given box;
 * if not, the atom is looped back the opposite wall of the
 * box
 *
 * @author Reid Van Lehn
 */
__device__ void inline validPos(double* pos, double boxSize) {
    if (*pos > boxSize) {
        *pos = *pos-boxSize;
    } else if (*pos < 0) {
        *pos = *pos+boxSize;
    } else {
        // value stays the same
    }
}



/**
 * @brief Updates positions of beads based on forces acting on them
 *
 * @todo Include stochastic operator to account for thermal fluctuations.
 *	Random numbers for the stochastic operator must match a normally
 *	distributed deviate with a mean 0 and variance 1
 * 
 */
__global__ void update_positions(double* bead_x, double* bead_y, double* bead_z, double* force_x, double* force_y, double* force_z) {
	int id = blockDim.x * blockIdx.x + threadIdx.x;
	//Update positions using dimensionaless stepsize as primary modifier.
	//Also include random stochastic operator
	bead_x[id] += (step_size*force_x[id] + sqrt(2*step_size));
	bead_y[id] += (step_size*force_y[id] + sqrt(2*step_size));
	bead_z[id] += (step_size*force_z[id] + sqrt(2*step_size));
	//Ensure positions stay within box
	validPos(bead_x+id, _boxX);
	validPos(bead_y+id, _boxY);
	/* REMOVED PERIODIC BC IN Z-DIRECTION */
	//r_z[i] = validPos(r_z[i], boxZ);

}



/**
 * @brief Initializes the simulation report output file
 *
 * This function writes the header for the simulation report
 */
void out_init(FILE* out_file) {
    /* Write number of beads for xyz plugin parsing */
    fprintf(out_file, "%d\n", num_beads);
    /* write parameters to file: num beads, steps, w_c, epsilon as comment */
    fprintf(out_file, "Lipids: %d  Steps: %d  w_c: %f  epsilon: %f"
            " Width: %f  Height %f\n", num_beads/3, steps,
            w_c_mult, epsilon, boxX, boxZ);
}



/**
 * @brief Writes a header for a simulation screenshot
 *
 * Need two lines between time steps
 * first one should be blank
 * second can be anything,
 * so write timestep number
 *
 */
void out_header(FILE* out_file, int timestep) {
	fprintf(out_file, "\nTimestep: %d\n", timestep);
}



/**
 * @brief Writes the data for a single simulation screenshot
 *
 * Outputs a single time step
 *
 */
void out_step(FILE* out_file, double* bead_x, double* bead_y, double* bead_z) {
	// Iterate through beads and write positions to file
	for (int i=0; i<num_beads; ++i) {
		if (i % 3)
			fprintf(out_file, "H %f %f %f\n", bead_x[i], bead_y[i], bead_z[i]);
        else
			fprintf(out_file, "O %f %f %f\n", bead_x[i], bead_y[i], bead_z[i]);
    }
}



int main()
{
	double* bead_x;
	double* bead_y;
	double* bead_z;
	double* force_x;
	double* force_y;
	double* force_z;
	double* _bead_x;
	double* _bead_y;
	double* _bead_z;
	double* _force_x;
	double* _force_y;
	double* _force_z;
	
	// allocate atoms property vectors
	bead_x = (double*)malloc(num_beads*sizeof(double));
	bead_y = (double*)malloc(num_beads*sizeof(double));
	bead_z = (double*)malloc(num_beads*sizeof(double));
	force_x = (double*)malloc(num_beads*sizeof(double));
	force_y = (double*)malloc(num_beads*sizeof(double));
	force_z = (double*)malloc(num_beads*sizeof(double));
	
	cudaMalloc((void**) &_bead_x, num_beads*sizeof(double));
	cudaMalloc((void**) &_bead_y, num_beads*sizeof(double));
	cudaMalloc((void**) &_bead_z, num_beads*sizeof(double));
	cudaMalloc((void**) &_force_x, num_beads*sizeof(double));
	cudaMalloc((void**) &_force_y, num_beads*sizeof(double));
	cudaMalloc((void**) &_force_z, num_beads*sizeof(double));
	
	// load the execution parameters
	init_params();
	
	// set up the bead arrangement
	init_positions(bead_x, bead_y, bead_z);
	cudaMemcpy(_bead_x, bead_x, num_beads*sizeof(double), cudaMemcpyHostToDevice);
	cudaMemcpy(_bead_y, bead_y, num_beads*sizeof(double), cudaMemcpyHostToDevice);
	cudaMemcpy(_bead_z, bead_z, num_beads*sizeof(double), cudaMemcpyHostToDevice);
	
	// init the output file
	FILE* out = fopen("out.txt", "w");
	out_init(out);
	
	for(int s=0; s < steps; ++s)
	{
		// clear previously calculated forces
		cudaMemset(_force_x, 0, num_beads*sizeof(double));
		cudaMemset(_force_y, 0, num_beads*sizeof(double));
		cudaMemset(_force_z, 0, num_beads*sizeof(double));
		
		// calculate the force constants for the first half of beads
		calc_forces_1st<<<num_beads/64, 32>>>(num_beads, _bead_x, _bead_y, _bead_z, _force_x, _force_y, _force_z);
		// calculate the force costants for the second half of beads
		calc_forces_2nd<<<num_beads/64, 32>>>(num_beads, _bead_x, _bead_y, _bead_z, _force_x, _force_y, _force_z);
		
		// update positions
		update_positions<<<num_beads/64, 64>>>(_bead_x, _bead_y, _bead_z, _force_x, _force_y, _force_z);
	}
	
	// fetch the result
	cudaMemcpy(_bead_x, bead_x, num_beads*sizeof(double), cudaMemcpyHostToDevice);
	cudaMemcpy(_bead_y, bead_y, num_beads*sizeof(double), cudaMemcpyHostToDevice);
	cudaMemcpy(_bead_z, bead_z, num_beads*sizeof(double), cudaMemcpyHostToDevice);
	
	// write the final status of the atoms
	out_header(out, steps);
	out_step(out, bead_x, bead_y, bead_z);
	
	// close output file
	fclose(out);
	
	// free atom properties vectors
	free(bead_x);
	free(bead_y);
	free(bead_z);
	free(force_x);
	free(force_y);
	free(force_z);
	
	cudaFree(_bead_x);
	cudaFree(_bead_y);
	cudaFree(_bead_z);
	cudaFree(_force_x);
	cudaFree(_force_y);
	cudaFree(_force_z);
}
